#include "ImageLoader.h"
#include "ImageFragment.h"
#include "../Common/ErrorCode.h"
#include "../Tools/randomGenerator.h"
#include "ML_IOTools.h"
#include "../Tools/ZLogManager.h"
#include "ImageFeature.h"
#include "ImageFeatureManager.h"

GLOBAL_NAMESPACE_BEGIN
NAMESPACE_MACHINE_LEARNING_BEGIN

class FragmentCmp
{
public:
    FragmentCmp(const ImageFragment& lhs)
    {
        lhs_ = &lhs;
    }
    bool operator()(const ImageFragment& rhs)
    {
        if (lhs_->ima_ == rhs.ima_ && 
            lhs_->params_.h == rhs.params_.h && lhs_->params_.w == rhs.params_.h &&
            lhs_->params_.u == rhs.params_.u && lhs_->params_.v == rhs.params_.v )
            return true;
        else
            return false;
    }

private:
    const ImageFragment* lhs_;
};

class FragmentSortByOrder
{
public:
    bool operator()(const ImageFragment& lhs, const ImageFragment& rhs) const
    {
        return lhs.getOrder() < rhs.getOrder();
    }
};

static double getMatCenterValue(const cv::Mat& mat)
{
    int x = mat.cols/2;
    int y = mat.rows/2;
    return mat.at<float>(x, y);
}

class ImageLoaderImpl
{
public:
    ImageLoaderImpl(DataManager* manager);
    friend class ImageLoader;

    void generateImageLoader();
    int  generateFragments(int num, const std::string& outFolder);
    int  extractFeature(ImageSpace::Image* ima, std::vector<double>& features, bool saveFeaImg);

    int getFragmentCount() const;
    const ImageFragment* getFragment(int idx) const;

    void addFragment(const ImageFragment& frag);
    void finalize();
    void loadFragments(const std::string& inFileName);
    void saveFragments(const std::string& outFileName);

private:
    void normalizedCrossRelation(cv::Mat& ret, const ImageSpace::Image* ima, const ImageFragment& fragment);
    void exponentiation(cv::Mat& ret, int exponent);
    void convolveDFT(const cv::Mat& A, const cv::Mat& B, cv::Mat& C);

private:
    DataManager *manager_;
    std::vector<cv::Mat*> images_;
    //std::vector<ImageFragmentParams> imageFragmentDictionary_;
    std::vector<ImageFragment> imageFragments_;
};

ImageLoaderImpl::ImageLoaderImpl(DataManager* manager)
    : manager_(manager)
{}

void ImageLoaderImpl::generateImageLoader()
{
    if (!manager_)
        return;

    int nCount = manager_->getImageCount();
    images_.clear();
    //images_.resize(nCount, NULL);
    for (int i = 0; i < nCount; i++)
    {
        ImageSpace::Image* image = manager_->getImageById(i);
        if (!image || image->getType() != ImageSpace::Image::IMAGE_Default)
            continue;
        //images_[i] = manager_->getImageById(i)->getData();
        images_.push_back(image->getData());
    }
}

// assume the images are all resizes to 32*32 size
int ImageLoaderImpl::generateFragments(int maxNum, const std::string& outFolder)
{
    int ret = ErrorCode::ERROR_NO_Error;
    if (!manager_ || manager_->getImageCount()<=0)
        return ErrorCode::ERROR_Invalid_Parameter;

    // build dictionary of fragments
    // only use square windows now
    const int minWindowSize = 4;
    const int maxWindowSize = 14;
    const int windowSize = 32;
    const int imageCount = manager_->getImageCount();
    Random_Generator::Start();
    while (imageFragments_.size()<maxNum)
    {
        int im = Random_Generator::getRandom(0, imageCount);
        int wh = Random_Generator::getRandom(minWindowSize, maxWindowSize+1);
        int u  = Random_Generator::getRandom(0, windowSize - wh);
        int v  = Random_Generator::getRandom(0, windowSize - wh);
        ImageSpace::Image* ima = manager_->getImageById(im);
        if (ima->getType() != ImageSpace::Image::IMAGE_Default)
            continue;
        if ( u+(int)(wh/2)+4>=windowSize || v+(int)(wh/2)+4>=windowSize 
            || u+(int)(wh/2)-3<0 || v+(int)(wh/2)-3<0 ) // avoid window is not totally in image
            continue;
        ImageFragmentParams param(u, v, wh, wh);
        param.setPath(outFolder.c_str());  // just for testing
        ImageFragment fragment(ima, param);
        if (std::find_if(imageFragments_.begin(), imageFragments_.end(), FragmentCmp(fragment))
            == imageFragments_.end())
            imageFragments_.push_back(fragment);
    }
    DebugLog << "ImageFragment number: " << imageFragments_.size();

    // generate the template
    int count = 0;
    for (std::vector<ImageFragment>::iterator it = imageFragments_.begin(); it!=imageFragments_.end(); it++)
    {
        it->finalize();
        it->setOrder(count++);
    }
    return ret;
}

int ImageLoaderImpl::extractFeature(ImageSpace::Image* ima, std::vector<double>& features, bool saveFeaImg)
{
    int ret = ErrorCode::ERROR_NO_Error;
    cv::Mat ncrMat;
    cv::Mat convolveMat;
    cv::Mat normalizedMat;
    std::vector<double> feature;
    ImageFeature* imgFeature = new ImageFeature(ima);
    for (int i = 0; i < imageFragments_.size(); i++)
    {
        normalizedCrossRelation(ncrMat, ima, imageFragments_.at(i));
        exponentiation(ncrMat, 10);
        convolveDFT(*imageFragments_.at(i).getMaskImage()->getData(), ncrMat, convolveMat);
        feature.push_back(getMatCenterValue(convolveMat));

        // for test
        if (saveFeaImg)
        {
            cv::normalize(convolveMat, normalizedMat, 0, 1, cv::NORM_MINMAX, -1, cv::Mat());
            imgFeature->addFeatureImage(normalizedMat*255);
        }
    }
    imgFeature->setFeatureVector(feature);
    ImageFeatureManager::getSingletonPtr()->addFeature(imgFeature);
    return ret;
}

void ImageLoaderImpl::normalizedCrossRelation(cv::Mat& ret, const ImageSpace::Image* ima, const ImageFragment& fragment)
{
    // assume the roi of the template are all ready
    cv::matchTemplate(*ima->getData(), *fragment.getTemplateImage()->getData(), ret, cv::TM_CCORR_NORMED);
    //DebugLog << ret;
    //cv::imwrite("ret0.jpg", ret*255);
    //cv::normalize(ret, ret, 0, 1, cv::NORM_MINMAX, -1, cv::Mat());
    //DebugLog << ret;
    //cv::imwrite("ret1.jpg", ret*255);
}

void ImageLoaderImpl::exponentiation(cv::Mat& mat, int p)
{
    cv::pow(mat, p, mat);
    //cv::normalize(mat, mat, 0, 1, cv::NORM_MINMAX, -1, cv::Mat());
    //DebugLog << mat;
    //cv::imwrite("ret2.jpg", mat*255);
}

void ImageLoaderImpl::convolveDFT(const cv::Mat& A, const cv::Mat& B, cv::Mat& C)
{
    // reallocate the output array if needed
    C.create(abs(A.rows-B.rows)+1, abs(A.cols-B.cols)+1, CV_32F);
    cv::Size dftSize;
    // calculate the size of DFT transform
    dftSize.width  = cv::getOptimalDFTSize(A.cols + B.cols - 1);
    dftSize.height = cv::getOptimalDFTSize(A.rows + B.rows - 1);

    cv::Mat dftA(dftSize, CV_32F, cv::Scalar(0));
    cv::Mat dftB(dftSize, CV_32F, cv::Scalar(0));
    cv::Mat tmp;

    cv::Mat roiA(dftA, cv::Rect(0, 0, A.cols, A.rows));
    cv::Mat roiB(dftB, cv::Rect(0, 0, B.cols, B.rows));
    A.convertTo(roiA, CV_32F);
    B.convertTo(roiB, CV_32F);

    //DebugLog << B;
    //DebugLog << dftB;
    //cv::imwrite("roiB.jpg", dftB);

    cv::dft(dftA, dftA, 0, A.rows);
    cv::dft(dftB, dftB, 0, B.rows);

    //DebugLog << dftB;
    //cv::imwrite("roiB1.jpg", dftB);

    cv::mulSpectrums(dftA, dftB, dftA, cv::DFT_COMPLEX_OUTPUT);

    //DebugLog << dftA;
    //cv::imwrite("roiA2.jpg", dftA);

    cv::dft(dftA, dftA, cv::DFT_INVERSE + cv::DFT_SCALE, C.rows);
    //cv::dft(dftA, dftA, cv::DFT_INVERSE, C.rows);

    //DebugLog << dftA;
    //cv::imwrite("roiA3.jpg", dftA);

    dftA(cv::Rect(0, 0, C.cols, C.rows)).copyTo(C);

    //cv::normalize(C, dftA, 0, 1, cv::NORM_MINMAX, -1, cv::Mat());
   // DebugLog << dftA;
    //cv::imwrite("C.jpg", dftA*255);
//    cv::imwrite("c.jpg", C);
    // all the temporary buffers will be deallocated automatically
}

void ImageLoaderImpl::saveFragments(const std::string& fileName)
{

}

int ImageLoaderImpl::getFragmentCount() const
{
    return imageFragments_.size();
}

const ImageFragment* ImageLoaderImpl::getFragment(int idx) const
{
    return &imageFragments_.at(idx);
}

void ImageLoaderImpl::addFragment(const ImageFragment& frag)
{
    imageFragments_.push_back(frag);
}

void ImageLoaderImpl::finalize()
{
    for (int i=0; i<imageFragments_.size(); i++)
    {
        imageFragments_.at(i).finalize();
    }
    std::sort(imageFragments_.begin(), imageFragments_.end(), FragmentSortByOrder());
}

//////////////////////////////////////////////////////////
///

ImageLoader::ImageLoader(DataManager* manager)
    : impl_(new ImageLoaderImpl(manager))
{
}

ImageLoader::~ImageLoader()
{
}

void ImageLoader::init()
{
    impl_->generateImageLoader();
}

int ImageLoader::extractFragments(int num, const std::string& outFolder)
{
    int ret = ErrorCode::ERROR_NO_Error;
    if (impl_->generateFragments(num, outFolder) != ErrorCode::ERROR_NO_Error)
        return ret;
    return ret;
}

int ImageLoader::extractFeaturesForImage(ImageSpace::Image* ima, std::vector<double>& features, bool saveFeaImg)
{
    return impl_->extractFeature(ima, features, saveFeaImg);
}

int ImageLoader::getFragmentCount() const
{
    return impl_->getFragmentCount();
}

const ImageFragment* ImageLoader::getFragment(int idx) const
{
    return impl_->getFragment(idx);
}

void ImageLoader::addFragment(const ImageFragment& frag)
{
    impl_->addFragment(frag);
}

void ImageLoader::finalize()
{
    impl_->finalize();
}

NAMESPACE_MACHINE_LEARNING_END
GLOBAL_NAMESPACE_END
